import os
import numpy as np
import nibabel as nib
import pandas as pd
from PIL import Image

# 输入路径
image_dir = "/data/suziren25/suziren/nnUNet/nnUNet_raw/abdominal_US/abdominal_US_organs/imagesTr/kidney"
label_dir = "/data/suziren25/suziren/nnUNet/nnUNet_raw/abdominal_US/abdominal_US_organs/labelsTr/kidney"

# 输出路径
output_root = "/data/suziren25/suziren/nnUNet/nnUNet_raw/Dataset506_KidneyUS"
os.makedirs(f"{output_root}/imagesTr", exist_ok=True)
os.makedirs(f"{output_root}/labelsTr", exist_ok=True)


def png_to_nifti(png_path, is_label=False):  # 🩵新增参数
    img = np.array(Image.open(png_path))
    if img.ndim == 3:
        img = img[..., 0]
    # 🩵 如果是 mask，则将 255→1，确保标签为 0/1
    if is_label:
        img = np.where(img == 255, 1, img)
    img = np.expand_dims(img, axis=-1)
    return nib.Nifti1Image(img.astype(np.uint8), affine=np.eye(4))


image_files = sorted([f for f in os.listdir(image_dir) if f.endswith(".png")])

mapping_data = []

for idx, fname in enumerate(image_files):
    case_id = f"case{idx + 1:03d}"
    image_path = os.path.join(image_dir, fname)
    label_path = os.path.join(label_dir, fname)

    if not os.path.exists(label_path):
        print(f"⚠️ Label not found for {fname}, skipping...")
        mapping_data.append({
            'original_image': fname,
            'original_label': fname,
            'case_id': case_id,
            'nnunet_image': f"{case_id}_0000.nii.gz",
            'nnunet_label': f"{case_id}.nii.gz",
            'status': 'skipped_no_label',
            'error': 'Label file not found'
        })
        continue

    try:
        # 转换图像
        img_nii = png_to_nifti(image_path, is_label=False)
        nib.save(img_nii, f"{output_root}/imagesTr/{case_id}_0000.nii.gz")

        # 转换mask（自动0/1规范化）
        mask_nii = png_to_nifti(label_path, is_label=True)  # 🩵传入 is_label=True
        nib.save(mask_nii, f"{output_root}/labelsTr/{case_id}.nii.gz")

        mapping_data.append({
            'original_image': fname,
            'original_label': fname,
            'case_id': case_id,
            'nnunet_image': f"{case_id}_0000.nii.gz",
            'nnunet_label': f"{case_id}.nii.gz",
            'status': 'success',
            'error': ''
        })

        print(f"✅ 转换成功: {fname} → {case_id}")

    except Exception as e:
        print(f"❌ 转换失败: {fname} - {e}")
        mapping_data.append({
            'original_image': fname,
            'original_label': fname,
            'case_id': case_id,
            'nnunet_image': f"{case_id}_0000.nii.gz",
            'nnunet_label': f"{case_id}.nii.gz",
            'status': 'failed',
            'error': str(e)
        })

# 保存映射关系为CSV
df = pd.DataFrame(mapping_data)
csv_path = os.path.join(output_root, "file_mapping.csv")
df.to_csv(csv_path, index=False, encoding='utf-8')

# 打印总结
success_count = len(df[df['status'] == 'success'])
failed_count = len(df[df['status'] == 'failed'])
skipped_count = len(df[df['status'] == 'skipped_no_label'])

print("\n" + "=" * 50)
print("📊 转换总结:")
print(f"✅ 成功转换: {success_count}")
print(f"❌ 失败: {failed_count}")
print(f"⚠️  跳过: {skipped_count}")
print(f"📝 映射文件: {csv_path}")
print("=" * 50)
